package com.jxauaia.contest.AI;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * 讯飞星火大模型WebSocket API客户端
 */
@Component
public class AIClient {

    private static final Logger logger = LoggerFactory.getLogger(AIClient.class);

    @Value("${ai.spark.appid:default_appid}")
    private String appId;

    @Value("${ai.spark.api-key:default_api_key}")
    private String apiKey;

    @Value("${ai.spark.api-secret:default_api_secret}")
    private String apiSecret;

    @Value("${ai.spark.timeout:30}")
    private int timeout;

    @Value("${ai.spark.api-url:wss://spark-api.xf-yun.com/v3.1/chat}")
    private String apiUrl;

    private final ObjectMapper objectMapper;

    public AIClient(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    /**
     * 发送问题到星火大模型并获取回答
     * @param question 用户问题
     * @param userId 用户ID，用于生成会话ID
     * @return AI回答
     */
    public String askQuestion(String question, String userId) {
        logger.info("开始处理AI问题，用户ID: {}", userId);

        try {
            // 生成会话ID
            String sessionId = "session_" + userId + "_" + System.currentTimeMillis();

            // 构建WebSocket URL，包含鉴权信息
            String url = buildAuthUrl();
            logger.debug("WebSocket URL: {}", url);

            // 创建请求体
            Map<String, Object> requestBody = createRequestBody(question, sessionId);
            String requestJson = objectMapper.writeValueAsString(requestBody);
            logger.debug("请求体: {}", requestJson);

            // 创建结果接收器
            CompletableFuture<String> resultFuture = new CompletableFuture<>();
            StringBuilder responseBuilder = new StringBuilder();

            // 创建WebSocket客户端
            WebSocketClient client = new WebSocketClient(new URI(url)) {
                @Override
                public void onOpen(ServerHandshake handshakeData) {
                    logger.info("WebSocket连接已打开，状态码: {}", handshakeData.getHttpStatus());
                    if (handshakeData.getHttpStatus() == 101) {
                        logger.info("WebSocket握手成功，发送请求");
                        send(requestJson);
                    } else {
                        logger.error("WebSocket握手失败，状态码: {}", handshakeData.getHttpStatus());
                        resultFuture.completeExceptionally(new RuntimeException("WebSocket握手失败，状态码: " + handshakeData.getHttpStatus()));
                        close();
                    }
                }

                @Override
                public void onClose(int code, String reason, boolean remote) {
                    logger.error("WebSocket连接已关闭: {} - {} - 远程关闭: {}", code, reason, remote);
                    if (!resultFuture.isDone()) {
                        if (code == 1002 && reason.contains("401")) {
                            // 使用文本块提升可读性
                            String errorMsg = """
                                    AI服务认证失败(401 Unauthorized)，可能原因：
                                    1. API密钥不正确或已过期
                                    2. 账户余额不足或调用次数已用完
                                    3. 讯飞平台账户状态异常

                                    请联系管理员检查API配置。
                                    """;
                            logger.error(errorMsg);
                            resultFuture.complete(errorMsg);
                        } else {
                            resultFuture.complete(responseBuilder.toString().isEmpty() ?
                                    "抱歉，AI服务暂时不可用，请稍后再试。错误代码: " + code : responseBuilder.toString());
                        }
                    }
                }

                @Override
                public void onMessage(String message) {
                    try {
                        logger.debug("收到消息: {}", message);
                        JsonNode response = objectMapper.readTree(message);

                        // 检查是否有错误
                        int code = response.path("header").path("code").asInt();
                        if (code != 0) {
                            String errorMsg = response.path("header").path("message").asText("Unknown error");
                            logger.error("API返回错误: {} - {}", code, errorMsg);
                            resultFuture.completeExceptionally(new RuntimeException("API错误: " + errorMsg));
                            close();
                            return;
                        }

                        // 获取状态
                        int status = response.path("header").path("status").asInt();

                        // 处理AI回复 - 修正路径及判断
                        if (response.has("payload") && response.path("payload").has("choices")) {
                            JsonNode choices = response.path("payload").path("choices");

                            // 检查不同版本API的返回格式
                            if (choices.has("text")) {
                                JsonNode textArray = choices.path("text");
                                if (textArray.isArray() && !textArray.isEmpty()) { // 替换size() > 0为isEmpty()
                                    String content = textArray.get(0).path("content").asText("");
                                    responseBuilder.append(content);
                                }
                            } else if (choices.isArray() && !choices.isEmpty()) { // 替换size() > 0为isEmpty()
                                String content = choices.get(0).path("content").path("text").asText("");
                                if (!content.isEmpty()) {
                                    responseBuilder.append(content);
                                }
                            }
                        }

                        // 如果是最后一条消息，完成结果
                        if (status == 2) {
                            resultFuture.complete(responseBuilder.toString());
                            close();
                        }
                    } catch (Exception e) {
                        logger.error("处理WebSocket消息时出错: {}", e.getMessage(), e);
                        resultFuture.completeExceptionally(e);
                        close();
                    }
                }

                @Override
                public void onError(Exception ex) {
                    logger.error("WebSocket连接错误: {}", ex.getMessage(), ex);
                    resultFuture.completeExceptionally(ex);
                    close();
                }
            };

            // 连接WebSocket
            client.connect();

            // 等待结果，设置超时
            try {
                String result = resultFuture.get(timeout, TimeUnit.SECONDS);
                logger.info("成功获取AI回答，用户ID: {}", userId);
                return result;
            } catch (TimeoutException e) {
                logger.error("获取AI回答超时: {}", e.getMessage(), e);
                client.close();
                return "抱歉，AI服务响应超时，请稍后再试。";
            } catch (ExecutionException e) {
                logger.error("获取AI回答出错: {}", e.getCause().getMessage(), e.getCause());
                client.close();
                return "抱歉，AI服务出现错误，请稍后再试。错误信息: " + e.getCause().getMessage();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                logger.error("获取AI回答被中断: {}", e.getMessage(), e);
                client.close();
                return "抱歉，获取AI回答被中断，请稍后再试。";
            }
        } catch (Exception e) {
            logger.error("处理AI问题时发生未预期错误: {}", e.getMessage(), e);
            return "抱歉，AI服务暂时不可用，请稍后再试。错误信息: " + e.getMessage();
        }
    }

    /**
     * 构建带鉴权信息的WebSocket URL - 完全按照官方示例实现
     */
    private String buildAuthUrl() throws Exception {
        logger.info("开始构建认证URL，API密钥: {}, API密钥长度: {}",
                apiKey.substring(0, 4) + "****", apiKey.length());

        // 解析URL获取host和path
        URL url = new URL(apiUrl.replace("wss://", "https://").replace("ws://", "http://"));
        String host = url.getHost();
        String path = url.getPath();

        // 获取当前时间
        SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        format.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = format.format(new Date());

        // 构建签名原始字符串
        String signatureOrigin = "host: " + host + "\n" +
                "date: " + date + "\n" +
                "GET " + path + " HTTP/1.1";

        logger.debug("签名原始字符串: {}", signatureOrigin);

        // 使用HMAC-SHA256算法计算签名 - 修正算法名称拼写
        Mac mac = Mac.getInstance("HmacSHA256"); // 修正为正确的算法名称
        SecretKeySpec secretKeySpec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "HmacSHA256");
        mac.init(secretKeySpec);
        byte[] signatureBytes = mac.doFinal(signatureOrigin.getBytes(StandardCharsets.UTF_8));
        String signature = Base64.getEncoder().encodeToString(signatureBytes);
        logger.debug("生成的签名: {}", signature);

        // 构建authorization - 修改为官方示例的格式
        String authorization = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"",
                apiKey, "hmac-sha256", "host date request-line", signature);
        logger.debug("生成的authorization: {}", authorization);

        // 构建URL参数 - 修改为官方示例的方式
        String httpUrl = "https://" + url.getHost() + url.getPath() +
                "?authorization=" + Base64.getEncoder().encodeToString(authorization.getBytes(StandardCharsets.UTF_8)) +
                "&date=" + date +
                "&host=" + host;

        // 转换为WebSocket URL
        String wsUrl = httpUrl.replace("https://", "wss://").replace("http://", "ws://");

        logger.debug("最终生成的URL: {}", wsUrl);
        return wsUrl;
    }

    /**
     * 创建请求体 - 按照官方示例格式
     */
    private Map<String, Object> createRequestBody(String question, String sessionId) {
        Map<String, Object> requestBody = new HashMap<>();

        // 构建header部分
        Map<String, Object> header = new HashMap<>();
        header.put("app_id", appId);
        header.put("uid", sessionId);
        requestBody.put("header", header);

        // 构建parameter部分 - 修正domain参数
        Map<String, Object> parameter = new HashMap<>();
        Map<String, Object> chat = new HashMap<>();

        // 根据API URL确定正确的domain参数
        String domain = "generalv3";
        if (apiUrl.contains("/v1.1/chat")) {
            domain = "lite";
        } else if (apiUrl.contains("/v3.1/chat")) {
            domain = "generalv3";
        } else if (apiUrl.contains("/chat/pro-128k")) {
            domain = "pro-128k";
        } else if (apiUrl.contains("/v3.5/chat")) {
            domain = "generalv3.5";
        } else if (apiUrl.contains("/chat/max-32k")) {
            domain = "max-32k";
        } else if (apiUrl.contains("/v4.0/chat")) {
            domain = "4.0Ultra";
        }

        chat.put("domain", domain);
        chat.put("temperature", 0.5);
        chat.put("max_tokens", 4096);
        parameter.put("chat", chat);
        requestBody.put("parameter", parameter);

        // 构建payload部分 - 修正格式
        Map<String, Object> payload = new HashMap<>();
        Map<String, Object> message = new HashMap<>();

        List<Map<String, String>> textList = new ArrayList<>();
        Map<String, String> userMessage = new HashMap<>();
        userMessage.put("role", "user");
        userMessage.put("content", question);
        textList.add(userMessage);

        message.put("text", textList);
        payload.put("message", message);
        requestBody.put("payload", payload);

        return requestBody;
    }
}